Personal To do

• Prettify benchmark results presentation
• Include benchmark for EMD
• Correct time measure for each backend (does time.perf_counter work on GPUs ??)
• Make real benchmark for Sinkhorn Knopp that is pretty for this PR (with every single backend and device)
• Put aforementioned benchmark in backend.py's docstring
• Check if set_gradient works as intended for tensorflow and add test
• Add to backend.py docstrings a comment about tensorflow requiring the numpy API
• Replace m2r2 by myst_parser (https://myst-parser.readthedocs.io/en/latest/)

Types of changes

• Docs change / refactoring / dependency upgrade
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)

Motivation and context / Related issue

• Backend for tensorflow.
• Also adds a benchmark folder which allows to easily benchmark every backend and present results in a html table (compatible with github comments and PR text).
• Also replaced m2r2 by myst_parser. It allows to use markdown files in our official docs. Myst_parser seems more maintained. It also happens to solve the issue with internal links introduced by the m2r2 update few PRs ago.

Note that tensorflow does not know how to transpose inplace, which can have a negative impact on performance for any function which requires the usage of a transpose (More information here: https://www.tensorflow.org/api_docs/python/tf/transpose#expandable-1)

Note that nothing was done regarding the device consistency of our calculation (unlike TorchBackend for instance). This is due to the fact that tensorflow chooses the device automatically without caring about the device where inputs are located. If the user wants a specific device, this should be done outside of POT, probably something like

with tf.device("/GPU:1"):
    ot.sinkhorn(args, **kwargs)

More information here: https://www.tensorflow.org/guide/gpu#overview

How has this been tested (if it applies)

Tested with Tensorflow 2.6.0 and 2.7.0. Benchmark with Tensorflow 2.6.0

Performance on Sinkhorn Knopp

The speedup is not as convincing as with cupy, might be due to how inefficient tensorflow transpose are. Still a great choice for very high dimensions.
Used pretty much the code from test/test_gpu.py::test_gpu_sinkhorn.
Tests were performed with Tensorflow 2.6.0, results in seconds, averaged over 100 runs.

python3 -m benchmarks.sinkhorn_knopp

Checklist

• The documentation is up-to-date with the changes I made.
• I have read the CONTRIBUTING document.
• All tests passed, and additional code has been covered with new tests.